from paddlenlp import Taskflow
import cv2
import numpy as np
import base64

encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), 95]
class NLP:
    def __init__(self):
        self.dialogue_model = Taskflow("dialogue", device_id=1)
        self.qa_model = Taskflow("question_answering", device_id=2)
        self.img_model = Taskflow("text_to_image", model="PaddlePaddle/disco_diffusion_ernie_vil-2.0-base-zh", device_id=3)
        # self.img_model = Taskflow("text_to_image", model="dalle-mini")

    def chat(self, text):
        result = self.dialogue_model([text])
        return result[0]

    
    def qa(self, text):
        result = self.qa_model(text)
        print("qa:", result)
        return result
    
    def img(self, text):
        text = [text]
        image_list = self.img_model(text)
        res = []
        for batch_index, batch_image in enumerate(image_list):
            for image_index_in_returned_images, each_image in enumerate(batch_image):
                cvImg = np.array(each_image)
                img = cv2.cvtColor(cvImg, cv2.COLOR_RGB2BGR)
                result, imgencode = cv2.imencode('.jpg', img, encode_param)
                data = np.array(imgencode)
                img = data.tobytes()
                # base64编码传输
                img = base64.b64encode(img).decode()
                res.append("data:image/jpg;base64," + img)
        return res
